#include "server.h"

EchoServer::EchoServer(int port, int maxEchoCount) 
    : port_(port), listenFd_(-1), epollFd_(-1), 
      totalConnections_(0), maxEchoCount_(maxEchoCount) {}

EchoServer::~EchoServer() {
    if (listenFd_ != -1) close(listenFd_);
    if (epollFd_ != -1)  close(epollFd_);
}

void EchoServer::setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags != -1) fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

int EchoServer::sendAll(int fd, const char* buffer, int len) {
    int total = 0;
    while (total < len) {
        int n = write(fd, buffer + total, len - total);
        if (n == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) return -1;
            return -2;
        }
        total += n;
    }
    return total;
}

bool EchoServer::init() {
    listenFd_ = socket(AF_INET, SOCK_STREAM, 0);
    if (listenFd_ == -1) {
        std::cerr << "Create socket failed: " << strerror(errno) << std::endl;
        return false;
    }

    int opt = 1;
    setsockopt(listenFd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(port_);
    if (bind(listenFd_, (struct sockaddr*)&addr, sizeof(addr)) == -1) {
        std::cerr << "Bind failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    if (listen(listenFd_, SOMAXCONN) == -1) {
        std::cerr << "Listen failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    epollFd_ = epoll_create1(0);
    if (epollFd_ == -1) {
        std::cerr << "Epoll create failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = listenFd_;
    epoll_ctl(epollFd_, EPOLL_CTL_ADD, listenFd_, &ev);

    std::cout << "Server initialized on port " << port_ 
              << " (max echo per connection: " << maxEchoCount_ << ")" << std::endl;
    return true;
}

void EchoServer::start() {
    struct epoll_event events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];

    while (true) {
        int nfds = epoll_wait(epollFd_, events, MAX_EVENTS, -1);
        if (nfds == -1) {
            if (errno == EINTR) continue;
            std::cerr << "Epoll wait failed: " << strerror(errno) << std::endl;
            break;
        }

        for (int i = 0; i < nfds; ++i) {
            int fd = events[i].data.fd;
            uint32_t event = events[i].events;

            if (fd == listenFd_) { // 新连接
                struct sockaddr_in clientAddr;
                socklen_t clientAddrLen = sizeof(clientAddr);
                int clientFd = accept(listenFd_, (struct sockaddr*)&clientAddr, &clientAddrLen);
                if (clientFd == -1) continue;

                setNonBlocking(clientFd);
                totalConnections_++;
                echoCount_[clientFd] = 0; // 初始化交互次数
                std::cout << "New connection (total: " << totalConnections_ 
                          << ") fd: " << clientFd << " from " 
                          << inet_ntoa(clientAddr.sin_addr) << ":" 
                          << ntohs(clientAddr.sin_port) << std::endl;

                struct epoll_event ev;
                ev.events = EPOLLIN;
                ev.data.fd = clientFd;
                epoll_ctl(epollFd_, EPOLL_CTL_ADD, clientFd, &ev);
            } 
            else { // 客户端事件
                if (event & EPOLLIN) { // 接收数据
                    int nbytes = read(fd, buffer, BUFFER_SIZE - 1);
                    if (nbytes <= 0) { // 连接关闭或错误
                        std::cout << "Connection closed (fd: " << fd 
                                  << ", total echoes: " << echoCount_[fd] << ")" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        echoCount_.erase(fd);
                        pendingData_.erase(fd);
                        continue;
                    }

                    buffer[nbytes] = '\0';
                    echoCount_[fd]++; // 递增交互次数
                    std::cout << "Server received (fd: " << fd 
                              << ", echo " << echoCount_[fd] << "/" << maxEchoCount_ 
                              << "): '" << buffer << "' (" << nbytes << " bytes)" << std::endl;

                    // 回显数据
                    int sendResult = sendAll(fd, buffer, nbytes);
                    if (sendResult == nbytes) {
                        std::cout << "Server echoed (fd: " << fd << "): '" << buffer << "'" << std::endl;

                        // 若达到最大交互次数，通知客户端关闭连接
                        if (echoCount_[fd] >= maxEchoCount_) {
                            std::cout << "Max echoes reached (fd: " << fd << "), closing connection" << std::endl;
                            epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                            close(fd);
                            echoCount_.erase(fd);
                        }
                    } 
                    else if (sendResult == -1) { // 缓存待发送数据
                        pendingData_[fd] = std::string(buffer, nbytes);
                        struct epoll_event ev;
                        ev.events = EPOLLIN | EPOLLOUT;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                    }
                    else { // 发送失败，关闭连接
                        std::cerr << "Send failed (fd: " << fd << ")" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        echoCount_.erase(fd);
                        pendingData_.erase(fd);
                    }
                }

                if (event & EPOLLOUT) { // 发送缓存数据
                    auto it = pendingData_.find(fd);
                    if (it == pendingData_.end()) {
                        struct epoll_event ev;
                        ev.events = EPOLLIN;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        continue;
                    }

                    std::string& data = it->second;
                    int sendResult = sendAll(fd, data.c_str(), data.size());
                    if (sendResult == data.size()) {
                        std::cout << "Server sent buffered data (fd: " << fd << "): '" << data << "'" << std::endl;
                        pendingData_.erase(it);
                        
                        // 检查是否达到最大交互次数
                        if (echoCount_[fd] >= maxEchoCount_) {
                            std::cout << "Max echoes reached (fd: " << fd << "), closing connection" << std::endl;
                            epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                            close(fd);
                            echoCount_.erase(fd);
                        } else {
                            struct epoll_event ev;
                            ev.events = EPOLLIN;
                            ev.data.fd = fd;
                            epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        }
                    } else {
                        std::cerr << "Buffered send failed (fd: " << fd << ")" << std::endl;
                        pendingData_.erase(it);
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        echoCount_.erase(fd);
                    }
                }
            }
        }
    }
}
